import math
import random
import struct
from typing import Any


def insert_random_character(s: str) -> str:
    """
    向 s 中下标为 pos 的位置插入一个随机 byte
    pos 为随机生成，范围为 [0, len(s)]
    插入的 byte 为随机生成，范围为 [32, 127]
    """
    # TODO
    if not s:
        return chr(random.randint(32, 127))
    
    pos = random.randint(0, len(s))
    random_char = chr(random.randint(32, 127))
    return s[:pos] + random_char + s[pos:]


def flip_random_bits(s: str) -> str:
    """
    基于 AFL 变异算法策略中的 bitflip 与 random havoc 实现相邻 N 位翻转（N = 1, 2, 4），其中 N 为随机生成
    从 s 中随机挑选一个 bit，将其与其后面 N - 1 位翻转（翻转即 0 -> 1; 1 -> 0）
    注意：不要越界
    """
    # TODO
    if not s:
        return s
    
    # 1. 选择N的值(1, 2或4)
    N = random.choice([1, 2, 4])
    
    # 2. 将字符串转换为bytearray以便进行位操作
    byte_arr = bytearray(s, 'utf-8')
    total_bits = len(byte_arr) * 8
    
    if total_bits == 0:
        return s
    
    # 3. 选择一个随机的起始bit位置
    start_bit = random.randint(0, total_bits - N)
    
    # 4. 翻转从start_bit开始的N个连续bit
    for i in range(N):
        if start_bit + i >= total_bits:
            break  # 边界检查
        
        # 计算字节位置和bit位置
        byte_pos = (start_bit + i) // 8
        bit_pos = (start_bit + i) % 8
        
        # 执行位翻转: XOR操作翻转特定位
        byte_arr[byte_pos] ^= (1 << bit_pos)
    
    return byte_arr.decode('utf-8',errors='replace')


def arithmetic_random_bytes(s: str) -> str:
    """
    基于 AFL 变异算法策略中的 arithmetic inc/dec 与 random havoc 实现相邻 N 字节随机增减（N = 1, 2, 4），其中 N 为随机生成
    字节随机增减：
        1. 取其中一个 byte，将其转换为数字 num1；
        2. 将 num1 加上一个 [-35, 35] 的随机数，得到 num2；
        3. 用 num2 所表示的 byte 替换该 byte
    从 s 中随机挑选一个 byte，将其与其后面 N - 1 个 bytes 进行字节随机增减
    注意：不要越界；如果出现单个字节在添加随机数之后，可以通过取模操作使该字节落在 [0, 255] 之间
    """
    if not s:
        return s
    
    # Choose N (1, 2, or 4)
    N = random.choice([1, 2, 4])
    
    byte_arr = bytearray(s, 'utf-8')
    length = len(byte_arr)
    
    if length == 0:
        return s
    # 添加长度检查
    if length < N:
        # 如果长度不足，只处理1字节
        N = 1
        if length == 0:
            return s  # 空字符串直接返回
    
    # Choose a random starting position
    start_pos = random.randint(0, length - N)
    
    # Modify N consecutive bytes
    for i in range(N):
        if start_pos + i >= length:
            break
            
        original_byte = byte_arr[start_pos + i]
        delta = random.randint(-35, 35)
        new_byte = (original_byte + delta) % 256
        byte_arr[start_pos + i] = new_byte
    
    return byte_arr.decode('utf-8',errors='replace')


def interesting_random_bytes(s: str) -> str:
    """
    基于 AFL 变异算法策略中的 interesting values 与 random havoc 实现相邻 N 字节随机替换为 interesting_value（N = 1, 2, 4），其中 N 为随机生成
    interesting_value 替换：
        1. 构建分别针对于 1, 2, 4 bytes 的 interesting_value 数组；
        2. 随机挑选 s 中相邻连续的 1, 2, 4 bytes，将其替换为相应 interesting_value 数组中的随机元素；
    注意：不要越界
    """
    # TODO
    if not s:
        return s
    
    # Interesting values for 1, 2, and 4 bytes
    interesting_1byte = [
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
        0x10, 0x20, 0x40, 0x7e, 0x7f, 0x80, 0x81, 0xc0, 0xfe, 0xff,
        ord('0'), ord('1'), ord('2'), ord('3'), ord('4'), ord('5'), ord('6'), ord('7'), ord('8'), ord('9'),
        ord('A'), ord('B'), ord('C'), ord('D'), ord('E'), ord('F'), ord('G'), ord('H'), ord('I'), ord('J'),
        ord('a'), ord('b'), ord('c'), ord('d'), ord('e'), ord('f'), ord('g'), ord('h'), ord('i'), ord('j')
    ]
    
    interesting_2byte = [
        0x0000, 0x0001, 0x0002, 0x0003, 0x0004, 0x0005, 0x0006, 0x0007, 0x0008, 0x0009, 0x000a, 0x000b,
        0x000c, 0x000d, 0x000e, 0x000f, 0x0010, 0x0020, 0x0040, 0x0080, 0x0100, 0x0200, 0x0400, 0x0800,
        0x1000, 0x2000, 0x4000, 0x8000, 0xffff, 0xfffe, 0xfffd, 0xfffc, 0x7fff, 0x3fff, 0x0fff, 0x07ff
    ]
    
    interesting_4byte = [
        0x00000000, 0x00000001, 0x00000002, 0x00000003, 0x00000004, 0x00000005, 0x00000006, 0x00000007,
        0x00000008, 0x00000009, 0x0000000a, 0x0000000b, 0x0000000c, 0x0000000d, 0x0000000e, 0x0000000f,
        0x00000010, 0x00000020, 0x00000040, 0x00000080, 0x00000100, 0x00000200, 0x00000400, 0x00000800,
        0x00001000, 0x00002000, 0x00004000, 0x00008000, 0x00010000, 0x00020000, 0x00040000, 0x00080000,
        0x00100000, 0x00200000, 0x00400000, 0x00800000, 0x01000000, 0x02000000, 0x04000000, 0x08000000,
        0x10000000, 0x20000000, 0x40000000, 0x80000000, 0xffffffff, 0xfffffffe, 0xfffffffd, 0xfffffffc,
        0x7fffffff, 0x3fffffff, 0x0fffffff, 0x07ffffff
    ]
    
    # Choose N (1, 2, or 4)
    N = random.choice([1, 2, 4])
    
    byte_arr = bytearray(s, 'utf-8')
    length = len(byte_arr)
    
    if length < N:
        return s  # Not enough bytes to replace
    
    # Choose a random starting position
    start_pos = random.randint(0, length - N)
    
    # Replace with interesting value
    if N == 1:
        interesting_val = random.choice(interesting_1byte)
        byte_arr[start_pos] = interesting_val % 256
    elif N == 2:
        interesting_val = random.choice(interesting_2byte)
        struct.pack_into('>H', byte_arr, start_pos, interesting_val)
    elif N == 4:
        interesting_val = random.choice(interesting_4byte)
        struct.pack_into('>I', byte_arr, start_pos, interesting_val)
    
    return byte_arr.decode('utf-8',errors='replace')


def havoc_random_insert(s: str):
    """
    基于 AFL 变异算法策略中的 random havoc 实现随机插入
    随机选取一个位置，插入一段的内容，其中 75% 的概率是插入原文中的任意一段随机长度的内容，25% 的概率是插入一段随机长度的 bytes
    """
    if not s:
        # If empty, just insert a random character
        return chr(random.randint(0, 255))
    
    insert_pos = random.randint(0, len(s))
    
    if random.random() < 0.75:
        # 75% chance: insert a substring from the original string
        if len(s) == 0:
            substring_length = 0
        else:
            substring_length = random.randint(1, min(16, len(s)))
        start = random.randint(0, len(s) - substring_length)
        substring = s[start:start+substring_length]
    else:
        # 25% chance: insert random bytes
        substring_length = random.randint(1, 16)
        substring = ''.join([chr(random.randint(0, 255)) for _ in range(substring_length)])
    
    return s[:insert_pos] + substring + s[insert_pos:]


def havoc_random_replace(s: str):
    """
    基于 AFL 变异算法策略中的 random havoc 实现随机替换
    随机选取一个位置，替换随后一段随机长度的内容，其中 75% 的概率是替换为原文中的任意一段随机长度的内容，25% 的概率是替换为一段随机长度的 bytes
    """
    # TODO
    if not s:
        return s
    
    # Choose a random position and length to replace
    replace_pos = random.randint(0, len(s) - 1)
    max_length = len(s) - replace_pos
    replace_length = random.randint(1, min(16, max_length))
    
    if random.random() < 0.75:
        # 75% chance: replace with a substring from the original string
        if len(s) == 0:
            substring_length = 0
        else:
            substring_length = random.randint(1, min(16, len(s)))
        start = random.randint(0, len(s) - substring_length)
        replacement = s[start:start+substring_length]
    else:
        # 25% chance: replace with random bytes
        replacement_length = random.randint(1, 16)
        replacement = ''.join([chr(random.randint(0, 255)) for _ in range(replacement_length)])
    
    # Ensure we don't go out of bounds
    actual_replace_length = min(replace_length, len(s) - replace_pos)
    return s[:replace_pos] + replacement[:actual_replace_length] + s[replace_pos+actual_replace_length:]

def mock_mutator(s:str):
    pass

def repeat_random_block(s: str) -> str:
    """
    随机选择一段字节块，并重复多次（2-4次）
    1. 随机选择起始位置 `pos` 和长度 `block_len`（确保不越界）
    2. 随机选择重复次数 `times`（2、3 或 4 次）
    3. 将该块重复插入到原位置
    """
    if not s:
        return s

    block_len = random.randint(1, min(8, len(s)))
    pos = random.randint(0, len(s) - block_len)
    block = s[pos : pos + block_len]
    times = random.randint(2, 4)

    return s[:pos] + (block * times) + s[pos + block_len:]

def xor_random_block(s: str) -> str:
    if not s:
        return s

    s_bytes = s.encode("utf-8")  # 或其他合适的编码
    block_len = random.randint(1, min(16, len(s_bytes)))
    pos = random.randint(0, len(s_bytes) - block_len)
    xor_mask = random.randint(1, 255)
    mutated_block = bytes([b ^ xor_mask for b in s_bytes[pos : pos + block_len]])
    result_bytes = s_bytes[:pos] + mutated_block + s_bytes[pos + block_len:]
    return result_bytes.decode("utf-8", errors="ignore")  # 忽略非法字符

class Mutator:

    def __init__(self) -> None:
        """Constructor"""
        self.mutators = [
            insert_random_character,
            flip_random_bits,
            arithmetic_random_bytes,
            interesting_random_bytes,
            havoc_random_insert,
            havoc_random_replace,
            repeat_random_block,
            xor_random_block,
            # mock_mutator
        ]

    def mutate(self, inp: Any) -> Any:
        mutator = random.choice(self.mutators)
        return mutator(inp)
